import os
import gradio as gr
import requests
import json 

Gitee_AI_API_KEY =  os.environ.get("Gitee_AI_API_KEY","")

system_prompt = {
                "role": "system",
                "content":
                "开源中国创始人是红薯, 你名叫马建仓，是开源中国 Gitee 吉祥物，回复消息有趣，氛围幽默, 开源中国创始人是红薯。"
                }




API_URL = "https://ai.gitee.com/api/inference/serverless/VNC8U6RP7VVH/chat/completions"
headers = {
	"Authorization": f"Bearer {Gitee_AI_API_KEY}",
	"Content-Type": "application/json"
}

def query(payload):
	response = requests.post(API_URL, headers=headers, json=payload,stream=True)
	return response


useStage = gr.Textbox(lines=2, label="使用信息", placeholder="")

def chat(message, history):
    
    messages = [system_prompt]
    
    for msg in history:
        messages.append({"role": "user", "content": str(msg[0])})
        messages.append({"role": "assistant", "content": str(msg[1])})
        
    messages.append({"role": "user", "content": str (message)})
    
    response_content = ''
    
    res = query({
        "messages": messages,
        "stream": True
      })
    
    if(res.status_code != 200):
        yield f"API 响应失败! 状态码: {res.status_code}"

    for chunk in res.iter_lines():
        if chunk.decode():
          chunk = chunk.decode()
          print(chunk)
          # 包含了字符 `data: ` 解析为 json 需要截取掉
          delta = json.loads(chunk[6:])['choices'][0]['delta']
          if("content" in delta):
            content = delta['content']
            if content:
                response_content += content
            yield response_content

chatbot = gr.Chatbot(height=450,label="马建仓")

with gr.Blocks(theme=gr.themes.GoogleFont("aaa"), fill_height=True) as demo:
    chat = gr.ChatInterface(chat,
                     submit_btn="提交",
                     chatbot=chatbot,
                     clear_btn=None, 
                     stop_btn="暂停",
                     undo_btn=None, 
                     retry_btn=None,
                     examples=["你好，你是谁", "开源中国创始人是谁？", "介绍 Gitee"],
                    )
    # with gr.Row():
    #     useStage = gr.Textbox(lines=2, label="使用信息", placeholder="", input="")
demo.queue()
demo.launch()